import os
import sys
import pandas as pd
from PyCmpltrtok.common import get_dir_name_ext


def split_csv(src_path, train=0.8, val=0.0, test=0.0):
    """
    读取CSV并切分

    :param src_path: 源CSV路径
    :param train: 训练集百分比
    :param val: 验证集百分比
    :param test: 测试集百分比
    :return: dict split_name-to-new_csv_path
    """
    # 读取
    df = pd.read_csv(src_path, delimiter=',', header=0)
    print(df[:10])

    # 洗牌
    df = df.sample(frac=1.0, random_state=666)
    print(df[:10])

    # 文件名
    xdir, xbasename, xext = get_dir_name_ext(src_path)
    path_train = os.path.join(xdir, f'{xbasename}-train.{xext}')
    path_val = os.path.join(xdir, f'{xbasename}-val.{xext}')
    path_test = os.path.join(xdir, f'{xbasename}-test.{xext}')
    for path in [path_train, path_val, path_test]:
        if os.path.exists(path):
            print(f'Path {path} already exists! Double check before splitting, please!', file=sys.stderr)
            return None

    # 切分
    name2path = {}
    xlen = len(df)
    n_train, n_val, n_test = 0, 0, 0
    if train:
        n_train = int(round(xlen * train))
        df_train = df.iloc[:n_train, :]
        name2path['train'] = path_train
        df_train.to_csv(path_train, index=False)
    if val:
        n_val = int(round(xlen * val))
        df_val = df.iloc[n_train:n_train+n_val, :]
        name2path['val'] = path_val
        df_val.to_csv(path_val, index=False)
    if test:
        df_test = df.iloc[n_train+n_val:, :]
        name2path['test'] = path_test
        df_test.to_csv(path_test, index=False)
    return name2path


if '__main__' == __name__:
    csv_path = r'D:\_dell7590_root\sync\1_usb\N1\large_sci.com.ai.large_data\NLP\waimai\waimai_10k.csv'
    result = split_csv(csv_path, 0.8, 0.1, 0.1)
    print(result)
